package com.example.websocket.service;

import com.example.websocket.ParamUtil;
import com.example.websocket.vo.WsMsgVo;
import io.micrometer.common.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * @author wwl
 * @version 1.0
 * @date 2024/8/26 14:26
 * @description:
 */
@Slf4j
public class Ws1Pool {

    private static final Map<String, WebSocketSession> pool = new ConcurrentHashMap<>();
    private static final Map<Long, List<String>> userMap = new ConcurrentHashMap<>();
    private static final ExecutorService threadPool = Executors.newFixedThreadPool(10);

    public static void add(WebSocketSession inbound) {
        pool.put(inbound.getId(), inbound);
        Map<String, String> map = ParamUtil.parser(inbound.getUri().getQuery());
        Long userId = Long.valueOf(map.get("userId"));
        log.info("add userId:{}", userId);
        List<String> lstInBound = userMap.computeIfAbsent(userId, k -> new ArrayList<>());
        lstInBound.add(inbound.getId());
        log.info("add connetion {},total size {}", inbound.getId(), pool.size());
    }

    public static void remove(WebSocketSession socket) {
        String sessionId = socket.getId();
        List<String> lstInBound = null;
        Map<String, String> map = ParamUtil.parser(socket.getUri().getQuery());
        Long userId = Long.valueOf(map.get("userId"));
        log.info("remove userId:{}", userId);
        if (StringUtils.isNotBlank(sessionId)) {
            lstInBound = userMap.get(userId);
            if (lstInBound != null) {
                lstInBound.remove(sessionId);
                if (lstInBound.isEmpty()) {
                    userMap.remove(userId);
                }
            }
        }

        pool.remove(sessionId);
        log.info("remove connetion {},total size {}", sessionId, pool.size());
    }

    /** 推送信息 */
    public static void broadcast(WsMsgVo vo) {
        Long userId = vo.getUserId();
        List<String> lstInBoundId;
        if (userId == null || userId == 0L) {
            // 发送给所有人
            lstInBoundId = userMap.values().stream().flatMap(List::stream).collect(Collectors.toList());
        } else {
            lstInBoundId = userMap.get(userId);
        }
        if (lstInBoundId == null || lstInBoundId.isEmpty()) {
            return;
        }
        threadPool.execute(() -> {
            try {
                for (String id : lstInBoundId) {
                    // 发送给指定用户
                    WebSocketSession connection = pool.get(id);
                    if (connection != null) {
                        synchronized (connection) {
                            TextMessage msg = new TextMessage(vo.getText());
                            connection.sendMessage(msg);
                        }
                    }
                }
            } catch (Exception e) {
                log.error("broadcast error: userId:{}", userId, e);
            }
        });
    }

}
